# File: fedentgate/utils/privacy_accounting.py
"""Advanced Privacy Accounting Tools for Differential Privacy"""

import numpy as np
from scipy.special import logsumexp

class PrivacyAccountant:
    def __init__(self, delta=1e-5):
        """
        Initialize privacy accountant
        Args:
            delta: Privacy parameter delta (default 1e-5)
        """
        self.delta = delta
        self.rdp_account = {}  # Stores RDP orders and accumulated values
        
    def add_rdp(self, sigma, sampling_rate, steps=1, orders=None):
        """
        Add RDP privacy cost for Gaussian mechanism
        Args:
            sigma: Noise multiplier (standard deviation)
            sampling_rate: Sampling probability q = L/N
            steps: Number of iterations
            orders: List of RDP orders to compute
        """
        if orders is None:
            # Default set of alpha values (orders)
            orders = [1 + x / 10 for x in range(1, 100)] + list(range(12, 64))
            
        for alpha in orders:
            # Compute RDP for Gaussian mechanism
            rdp_val = alpha / (2 * sigma**2)
            # Account for subsampling amplification
            rdp_subsampled = self._compute_rdp_subsampling(rdp_val, sampling_rate, alpha)
            # Accumulate over steps
            if alpha in self.rdp_account:
                self.rdp_account[alpha] += rdp_subsampled * steps
            else:
                self.rdp_account[alpha] = rdp_subsampled * steps
                
    def _compute_rdp_subsampling(self, rdp, q, alpha):
        """
        Compute RDP for subsampled Gaussian mechanism
        Args:
            rdp: RDP value without subsampling
            q: Sampling probability
            alpha: RDP order
        Returns:
            rdp_subsampled: Amplified RDP value
        """
        # Using the bound from Wang et al. (ICML 2019)
        if q == 0:
            return 0
            
        # Compute log terms safely
        log_term1 = np.log(1 - q)
        log_term2 = np.log(q) + rdp + np.log1p(-np.exp(-rdp))
        
        # Combine using logsumexp for numerical stability
        log_terms = [log_term1, log_term2]
        rdp_subsampled = (1 / (alpha - 1)) * logsumexp(log_terms)
        
        return rdp_subsampled
    
    def get_epsilon(self, delta=None):
        """
        Convert accumulated RDP to (epsilon, delta)-DP
        Args:
            delta: Target delta (uses instance delta if None)
        Returns:
            epsilon: Computed privacy budget
        """
        if delta is None:
            delta = self.delta
            
        if not self.rdp_account:
            return float('inf')
            
        # Find minimum epsilon over all orders
        epsilons = []
        for alpha, rdp_sum in self.rdp_account.items():
            epsilon = rdp_sum + (np.log(1 / delta) / (alpha - 1))
            epsilons.append(epsilon)
            
        return min(epsilons)
    
    def get_privacy_spent(self):
        """
        Get current privacy expenditure
        Returns:
            (epsilon, delta): Current privacy parameters
        """
        epsilon = self.get_epsilon()
        return epsilon, self.delta
    
    def reset(self):
        """Reset accumulated privacy costs"""
        self.rdp_account = {}
